import torch
import torch.nn.functional as F


def structure_loss(pred, mask):
    '''
    1,1,448,448
    边界像素接收更大的权重
    '''
    weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)
    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))

    pred = torch.sigmoid(pred)
    inter = ((pred * mask)*weit).sum(dim=(2, 3))
    union = ((pred + mask)*weit).sum(dim=(2, 3))
    wiou = 1 - (inter + 1)/(union - inter+1)
    return (wbce + wiou).mean()
def dice_loss(pred, mask):

    # pred = pred.squeeze(dim = 1)

    smooth = 1

    # dice系数的定义
    dice = 2 * (pred * mask).sum(dim = 1).sum(dim = 1).sum(dim = 1) / (
                pred.pow(2).sum(dim = 1).sum(dim = 1).sum(dim = 1) +
                mask.pow(2).sum(dim = 1).sum(dim = 1).sum(dim = 1) + smooth)

    # 返回的是dice距离
    return torch.clamp((1 - dice).mean(), 0, 1)

def tversky_loss(pred, mask):
    smooth = 1

    # dice系数的定义
    dice = (pred * mask).sum(dim=1).sum(dim=1).sum(dim=1) / ((pred * mask).sum(dim=1).sum(dim=1).sum(dim=1)+
                                        0.3 * (pred * (1 - mask)).sum(dim=1).sum(dim=1).sum(dim=1) + 0.7 * ((1 - pred) * mask).sum(dim=1).sum(dim=1).sum(dim=1) + smooth)

    # 返回的是dice距离
    return torch.clamp((1 - dice).mean(), 0, 2)
def get_loss():
    return structure_loss

if __name__ == '__main__':
    import torch
    a=torch.zeros((1,1,8,8))
    b = torch.ones((1, 1, 8, 8))
    loss=get_loss()
    print(loss(a,b))